import pandas as pd
import numpy as np
import logging
from typing import Union, Dict
from .exceptions import IndicatorError

logger = logging.getLogger(__name__)

def calculate_sma(series: pd.Series, window: int, return_series: bool = False) -> Union[float, pd.Series]:
    """计算简单移动平均线 (SMA)"""
    try:
        if not isinstance(series, pd.Series):
            raise IndicatorError("Input must be a pandas Series")
        
        if window <= 0:
            raise IndicatorError("Window must be positive")
        
        if len(series) < window:
            logger.warning(f"Series length ({len(series)}) is less than window ({window})")
            if return_series:
                return pd.Series(index=series.index, dtype=float)
            else:
                return np.nan
        
        sma_series = series.rolling(window=window, min_periods=1).mean()
        
        if return_series:
            return sma_series
        else:
            return sma_series.iloc[-1] if not sma_series.empty else np.nan
            
    except Exception as e:
        logger.error(f"Error calculating SMA: {str(e)}")
        raise IndicatorError(f"SMA calculation failed: {str(e)}")

def calculate_rsi(series: pd.Series, window: int = 14, return_series: bool = False) -> Union[float, pd.Series]:
    """计算相对强弱指数 (RSI)"""
    try:
        if not isinstance(series, pd.Series):
            raise IndicatorError("Input must be a pandas Series")
        
        if window <= 0:
            raise IndicatorError("Window must be positive")
        
        if len(series) < window + 1:
            logger.warning(f"Series length ({len(series)}) is insufficient for RSI calculation (need {window + 1})")
            if return_series:
                return pd.Series(index=series.index, dtype=float)
            else:
                return np.nan
        
        delta = series.diff()
        gain = delta.where(delta > 0, 0).rolling(window=window, min_periods=1).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=window, min_periods=1).mean()
        
        # 处理除零情况
        rs = gain / loss
        rs = rs.replace([np.inf, -np.inf], np.nan)
        
        rsi_series = 100 - (100 / (1 + rs))
        
        # 处理边界情况
        rsi_series = rsi_series.fillna(50)  # 默认值为50
        rsi_series = rsi_series.clip(0, 100)  # 限制在0-100范围内
        
        if return_series:
            return rsi_series
        else:
            return rsi_series.iloc[-1] if not rsi_series.empty else 50.0
            
    except Exception as e:
        logger.error(f"Error calculating RSI: {str(e)}")
        raise IndicatorError(f"RSI calculation failed: {str(e)}")

def calculate_atr(high: pd.Series, low: pd.Series, close: pd.Series, 
                  window: int = 14, return_series: bool = False) -> Union[float, pd.Series]:
    """计算平均真实波幅 (ATR)"""
    try:
        # 验证输入
        if not all(isinstance(s, pd.Series) for s in [high, low, close]):
            raise IndicatorError("All inputs must be pandas Series")
        
        if window <= 0:
            raise IndicatorError("Window must be positive")
        
        if len(high) != len(low) or len(low) != len(close):
            raise IndicatorError("All series must have the same length")
        
        if len(high) < window:
            logger.warning(f"Series length ({len(high)}) is less than window ({window})")
            if return_series:
                return pd.Series(index=high.index, dtype=float)
            else:
                return np.nan
        
        # 计算真实波幅
        tr_data = pd.DataFrame({
            'hl': high - low,
            'hc': abs(high - close.shift(1)),
            'lc': abs(low - close.shift(1))
        })
        
        tr = tr_data.max(axis=1)
        atr_series = tr.rolling(window=window, min_periods=1).mean()
        
        if return_series:
            return atr_series
        else:
            return atr_series.iloc[-1] if not atr_series.empty else np.nan
            
    except Exception as e:
        logger.error(f"Error calculating ATR: {str(e)}")
        raise IndicatorError(f"ATR calculation failed: {str(e)}")

def calculate_bollinger_bands(series: pd.Series, window: int = 20, num_std: float = 2.0) -> Dict[str, pd.Series]:
    """计算布林格带"""
    try:
        if not isinstance(series, pd.Series):
            raise IndicatorError("Input must be a pandas Series")
        
        if window <= 0:
            raise IndicatorError("Window must be positive")
        
        if num_std <= 0:
            raise IndicatorError("Number of standard deviations must be positive")
        
        sma = series.rolling(window=window).mean()
        std = series.rolling(window=window).std()
        
        upper_band = sma + (std * num_std)
        lower_band = sma - (std * num_std)
        
        return {
            'middle': sma,
            'upper': upper_band,
            'lower': lower_band,
            'bandwidth': (upper_band - lower_band) / sma
        }
        
    except Exception as e:
        logger.error(f"Error calculating Bollinger Bands: {str(e)}")
        raise IndicatorError(f"Bollinger Bands calculation failed: {str(e)}")